import re
from typing import List, Optional

from instruction_re.core.datatypes import Preffix, Relation


class PredictionTripleFormatter:
    """
    Turns raw Model output into NER spans (start_idx, end_idx, label)
    """

    # answer_templates = ["is an", "is a"]  # TODO move this (get rid of literals)

    def format_answer_triples(
        self, context: str, prediction: str, options: List[str]
    ) -> List[Relation]:
        """
        Based on model prediction and context create entity spans
        :param options:
        :param context:
        :param prediction:
        :return:
        """

        triples = []
        source_sentence = context.replace(
            Preffix.CONTEXT.value,
            "",
            1,  # replace only the first occurrence of substring
        )
        
        
        
        # prediction = prediction.replace(
        #     Preffix.ANSWER.value,
        #     "",
        #     1,  # replace only the first occurrence of substring
        # )

        prediction = prediction.strip(
            "."
        )  # Because answer in train data always ends with '.'
        prediction_parts = prediction.split(",")

        for prediction_part in prediction_parts:
            spans = self._get_triple_from_part(
                prediction_part, source_sentence, options
            )
            if spans is None:
                continue

            spans = [span for span in spans if span.type in options]
            triples.extend(spans)

        return triples

    def _get_triple_from_part(
        self, prediction_part: str, source_sentence: str, options: List[str]
    ) -> Optional[List[Relation]]:
        """
        Gets entity span from part of prediction
        :param prediction_part: Olga is a PER
        :param source_sentence: Today Olga decided to sleep a lot.
        :return: (sub, type, obj)
        """

        if not any([template in prediction_part for template in options]):
            print("error pred：", prediction_part)
            return None

        for answer_template in options:
            _prediction_part = prediction_part.split(answer_template, maxsplit=2)

            if len(_prediction_part) != 2:
                continue

            sub, obj = (
                _prediction_part[0],
                _prediction_part[1],
            )

            sub = sub.strip(" ").rstrip(" ")
            rel_type = answer_template.strip(" ").rstrip(" ")
            obj = obj.strip(" ").rstrip(" ")

            # 检查 sub 和 obj 是否在 source_sentence 中
            try:
                matches = list(re.finditer(sub, source_sentence))
            except re.error:  # unbalanced parenthesis at position
                return None

            if len(matches) == 0:
                return None

            try:
                matches = list(re.finditer(obj, source_sentence))
            except re.error:  # unbalanced parenthesis at position
                return None

            if len(matches) == 0:
                return None

            rel = Relation(head=sub, tail=obj, type=rel_type)

            return [rel]

        return None


if __name__ == "__main__":
    example = {
        "sentence": "John Wilkes Booth , who assassinated President Lincoln.",
        "output": "John Wilkes Booth kill President Lincoln.",
    }

    options = [
        "organization based in",
        "located in",
        "live in",
        "work for",
        "kill",
    ]

    formatter = PredictionTripleFormatter()

    print(
        formatter.format_answer_triples(
            context=example["sentence"],
            prediction=example["output"],
            options=options,
        )
    )
